#include "icons.h"

IconFileList Icons::ExtractIcons(TCHAR* filename) {
  IconFileList icon_list = {0};
  icon_list.icon_count = ExtractIconEx(filename, -1, NULL, NULL, 0);
  _tcscpy(icon_list.iconfile, filename);
  if (icon_list.icon_count > 0) {
    icon_list.icon_data_list = new IconFileData[icon_list.icon_count];
    if (!icon_list.icon_data_list)
      return icon_list;
  }

  HDC hdc = CreateCompatibleDC(NULL);
  for (int index = 0; index <  icon_list.icon_count; index++) {
    HICON hicon = ExtractIcon(GetModuleHandle(NULL), filename, index);
    if (hicon) {
      ICONINFO info = {0};
      GetIconInfo(hicon, &info);
      if (!info.fIcon) {
        DestroyIcon(hicon);
        continue;
      }

      memset(&icon_list.icon_data_list[index], 0 , sizeof(IconFileData));
      BITMAPINFO color_bitmap = {sizeof(BITMAPINFOHEADER)};
      GetDIBits(hdc, info.hbmColor, 0, 0, NULL, &color_bitmap, DIB_RGB_COLORS);
      BITMAPINFO mono_bitmap = {sizeof(BITMAPINFOHEADER)};
      GetDIBits(hdc, info.hbmMask, 0, 0, NULL, &mono_bitmap, DIB_RGB_COLORS);
      icon_list.icon_data_list[index].dir.count = 1;
      icon_list.icon_data_list[index].dir.reserved = 0;
      icon_list.icon_data_list[index].dir.type = 1;
      icon_list.icon_data_list[index].dir.entries[0].height = 
          color_bitmap.bmiHeader.biHeight;
      icon_list.icon_data_list[index].dir.entries[0].width = 
          color_bitmap.bmiHeader.biWidth;
      icon_list.icon_data_list[index].dir.entries[0].reserved = 0;
      icon_list.icon_data_list[index].dir.entries[0].planes = 
          color_bitmap.bmiHeader.biPlanes;
      icon_list.icon_data_list[index].dir.entries[0].bit_count = 
          color_bitmap.bmiHeader.biBitCount;
      if (icon_list.icon_data_list[index].dir.entries[0].bit_count >= 8) {
        icon_list.icon_data_list[index].dir.entries[0].color_count = 0;
      } else {
        icon_list.icon_data_list[index].dir.entries[0].color_count =
            icon_list.icon_data_list[index].dir.entries[0].bit_count * 
            icon_list.icon_data_list[index].dir.entries[0].planes;
      }
      icon_list.icon_data_list[index].dir.entries[0].image_offset = 
          sizeof(IconDir);

      int color_info_size = sizeof(BITMAPINFOHEADER); 
      if (color_bitmap.bmiHeader.biBitCount > 8)   { 
        if (color_bitmap.bmiHeader.biCompression & BI_BITFIELDS) 
          color_info_size += sizeof(RGBQUAD) * 3; 
      } else {
        color_info_size += sizeof(RGBQUAD) * (
          1 << color_bitmap.bmiHeader.biBitCount);
      }
      int mask_info_size = sizeof(BITMAPINFOHEADER); 
      if (mono_bitmap.bmiHeader.biBitCount > 8)   { 
        if (mono_bitmap.bmiHeader.biCompression & BI_BITFIELDS) 
          mask_info_size += sizeof(RGBQUAD) * 3; 
      } else {
        mask_info_size += sizeof(RGBQUAD) * (
          1 << mono_bitmap.bmiHeader.biBitCount);
      }
      BITMAPINFO* color_info = (BITMAPINFO*)malloc(color_info_size);
      BITMAPINFO* mono_info = (BITMAPINFO*)malloc(mask_info_size);
      memcpy(color_info, &color_bitmap, sizeof(color_bitmap));
      memcpy(mono_info, &mono_bitmap, sizeof(mono_bitmap));
      icon_list.icon_data_list[index].dir.entries[0].bytes_in_res = 
          color_info_size + color_bitmap.bmiHeader.biSizeImage + 
          mono_bitmap.bmiHeader.biSizeImage;
      icon_list.icon_data_list[index].icon_data = (byte*)malloc(
          icon_list.icon_data_list[index].dir.entries[0].bytes_in_res +
          color_info_size);
      if (!icon_list.icon_data_list[index].icon_data) {
        DestroyIcon(hicon);
        continue;
      }

      GetDIBits(hdc, info.hbmColor, 0, color_bitmap.bmiHeader.biHeight, 
                icon_list.icon_data_list[index].icon_data + 
                color_info_size, color_info, DIB_RGB_COLORS);
      GetDIBits(hdc, info.hbmMask, 0, mono_bitmap.bmiHeader.biHeight, 
                icon_list.icon_data_list[index].icon_data + 
                color_info_size + color_bitmap.bmiHeader.biSizeImage,
                mono_info, DIB_RGB_COLORS);

      color_info->bmiHeader.biHeight *= 2;
      memcpy(icon_list.icon_data_list[index].icon_data, color_info, 
             color_info_size);

      free(color_info);
      free(mono_info);

      DestroyIcon(hicon);
    }
  }
  DeleteDC(hdc);
  return icon_list;
}

void Icons::FreeData(IconFileList& iconlist) {
  for (int index = 0; index < iconlist.icon_count; index++) {
    if (iconlist.icon_data_list[index].icon_data) {
      free(iconlist.icon_data_list[index].icon_data);
      iconlist.icon_data_list[index].icon_data = NULL;
    }
  }
}

BOOL CALLBACK MyEnumProcedure(HANDLE moudle, LPCTSTR type, 
                              LPTSTR name, LONG param)	{
  LPTSTR resource_id;
  IconFile* iconfile = (IconFile*)param;
  HMODULE library = iconfile->handle;

  resource_id = name;

  if (type == RT_GROUP_ICON) {
    if (iconfile->index-- == 0){
      HRSRC group_rsrc = FindResource(library, resource_id, RT_GROUP_ICON);
      if (group_rsrc == NULL) {
        return TRUE;
      }
      HGLOBAL group_global = LoadResource(library, group_rsrc);
      if (group_global == NULL) {
        return TRUE;
      }

      GrpIconDir* icon_dir = (GrpIconDir*)LockResource(group_global);
      if (icon_dir == NULL) {
        return TRUE;
      }

      DWORD size = sizeof(IconDir);
      for (int index = 0; index < icon_dir->count; index++) {
        size += index * sizeof(IconDirEntry) + 
            icon_dir->entries[index].bytes_in_res;
      }
      iconfile->datasize = size;
      iconfile->data = (byte*)malloc(size);
      if (!iconfile->data)
        return TRUE;

      IconDir* filedir = (IconDir*)iconfile->data;
      filedir->count = icon_dir->count;
      filedir->reserved = 0;
      filedir->type = 1;
      int offset = sizeof(IconDir) + 
          (icon_dir->count - 1) * sizeof(IconDirEntry);
      for (int index = 0; index < icon_dir->count; index++) {
        filedir->entries[index].width = icon_dir->entries[index].width;
        filedir->entries[index].height = icon_dir->entries[index].height;
        filedir->entries[index].color_count = 
            icon_dir->entries[index].color_count;
        filedir->entries[index].reserved = icon_dir->entries[index].reserved;
        filedir->entries[index].planes = icon_dir->entries[index].planes;
        filedir->entries[index].bit_count = icon_dir->entries[index].bit_count;
        filedir->entries[index].bytes_in_res = 
            icon_dir->entries[index].bytes_in_res;
        if (index > 0)
          offset += icon_dir->entries[index - 1].bytes_in_res;
        filedir->entries[index].image_offset = offset;
        HRSRC image_rsrc = FindResource(
            library, MAKEINTRESOURCE(icon_dir->entries[index].id), RT_ICON);
        HGLOBAL image_global= LoadResource(library, image_rsrc);
        IconImage* iconimage = (IconImage*)LockResource(image_global);
        int size = SizeofResource(library, image_rsrc);
        memcpy((byte*)filedir + offset, iconimage, 
                icon_dir->entries[index].bytes_in_res);
      }
    }
  }

  return TRUE;
}

IconFile Icons::ExtractExeIcon(TCHAR* filename, int index) {
  IconFile iconfile = {0};
  _tcscpy(iconfile.exefile, filename);
  iconfile.index = index;
  HMODULE library = LoadLibraryEx(filename, NULL, LOAD_LIBRARY_AS_DATAFILE);
  if (!library)
    return iconfile;

  iconfile.handle = library;

  EnumResourceNames(library, RT_GROUP_ICON, 
                    (ENUMRESNAMEPROC)MyEnumProcedure, (LPARAM)&iconfile);

  FreeLibrary(library);
  return iconfile;
}

void Icons::FreeData(IconFile& iconfile) {
  if (iconfile.datasize > 0 && iconfile.data) {
    free(iconfile.data);
    iconfile.data = 0;
  }
}